import torch
from torch_sparse import SparseTensor as st, diag 
from torch_sparse.matmul import matmul as spspmm

class JacobPreconditioner:
    
    def __init__(self,indices,values,m) -> None:
        
        A=st(row=indices[0],col=indices[1],value=values,sparse_sizes=(m,m),is_sorted=True).to_symmetric()

        idx=torch.arange(m,dtype=indices.dtype,device=indices.device)
        self.__diag=st(idx,None,idx,1/diag.get_diag(A))

    def apply(self,indices,values,m,n):
        B=st(row=indices[0],col=indices[1],value=values,sparse_sizes=(m,n),is_sorted=True)

        C=spspmm(self.__diag,B)
        row, col, value = C.coo()
        return torch.stack([row, col], dim=0), value


class ColNormPreconditioner:

    """
    order is tha same as torch.linalg.norm(...,ord)

    for column absoluate sum preconditioner order=0
    for column Euclid preconditioner order=2
    for column max preconditioner order=float('inf')


    """
    def __init__(self,indices,values,m,order) -> None:
        A=st(row=indices[0],col=indices[1],value=values,sparse_sizes=(m,m),is_sorted=True)
        col_ptr,_,val=A.csc()

        d=torch.empty(m,dtype=values.dtype,device=values.device)

        for i in range(m):
            d=torch.linalg.norm(val[col_ptr[i]:col_ptr[i+1]].abs(),order)

        idx=torch.arange(m,dtype=indices.dtype,device=indices.device)
        self.__diag=st(idx,None,idx,1/d)

    def apply(self,indices,values,m,n):
        B=st(row=indices[0],col=indices[1],value=values,sparse_sizes=(m,n),is_sorted=True)

        C=spspmm(self.__diag,B)
        row, col, value = C.coo()
        return torch.stack([row, col], dim=0), value






